use core::panic;
use std::path::PathBuf;

use crate::data::{
    transfer::ImageMeta,
    utils::{read_classes_file, read_label_file},
};
use anyhow::Result;
use burn::{
    config::Config,
    data::dataset::{
        Dataset as BurnDataset,
        transform::ComposedDataset,
        vision::{Annotation, ImageDatasetItem, ImageFolderDataset, PixelDepth},
    },
    serde::Serialize,
};
use hdf5::{Dataset, File as H5File, types::VarLenUnicode};
use serde::Deserialize;

pub struct HWDBDataset {
    pub train: ComposedDataset<ImageFolderDataset>,
    pub test: ComposedDataset<ImageFolderDataset>,
}

impl HWDBDataset {
    pub fn new(
        train_dir_vec: &Vec<String>,
        test_dir_vec: &Vec<String>,
        classes_filepath: &str,
        label_filename: &str,
    ) -> anyhow::Result<HWDBDataset> {
        let classes_path = PathBuf::from(classes_filepath);
        assert!(classes_path.exists(), "classes file path is not exist");
        let Some(classes_path) = classes_path.to_str() else {
            panic!("can't find classes file")
        };
        let classes = read_classes_file(classes_path)?;

        let mut train_datasets = Vec::new();
        let mut test_datasets = Vec::new();

        for item in train_dir_vec {
            println!("build train dataset {}", item);
            let subdir = PathBuf::from(item);

            assert!(
                subdir.exists(),
                "train sub dir is not exist {}",
                subdir.display()
            );

            let label_file = subdir.join(label_filename);
            assert!(
                label_file.exists(),
                "label file is not exist {}",
                label_file.display()
            );

            let Some(label_file) = label_file.to_str() else {
                panic!("can't find label file")
            };
            let label_file_contents = read_label_file(label_file)?;
            let mut labels = Vec::new();
            for (path, label) in label_file_contents {
                let image_path = subdir.join(path);
                labels.push((image_path, label));
            }

            train_datasets.push(ImageFolderDataset::new_classification_with_items(
                labels, &classes,
            )?);
        }

        for item in test_dir_vec {
            println!("build test dataset {}", item);
            let subdir = PathBuf::from(item);
            assert!(
                subdir.exists(),
                "train sub dir is not exist {}",
                subdir.display()
            );

            let label_file = subdir.join(label_filename);
            assert!(
                label_file.exists(),
                "label file is not exist {}",
                label_file.display()
            );

            let Some(label_file) = label_file.to_str() else {
                panic!("can't find label file")
            };
            let label_file_contents = read_label_file(label_file)?;
            let mut labels = Vec::new();
            for (path, label) in label_file_contents {
                let image_path = subdir.join(path);
                labels.push((image_path, label));
            }

            test_datasets.push(ImageFolderDataset::new_classification_with_items(
                labels, &classes,
            )?);
        }
        Ok(HWDBDataset {
            train: ComposedDataset::new(train_datasets),
            test: ComposedDataset::new(test_datasets),
        })
    }
}

#[derive(Config, Debug)]
pub struct HWDBDatasetConfig {
    pub trn_dir_vec: Vec<String>,
    pub tst_dir_vec: Vec<String>,
    pub classes_filepath: String,
    pub label_filename: String,
}

impl HWDBDatasetConfig {
    pub fn init(&self) -> HWDBDataset {
        HWDBDataset::new(
            &self.trn_dir_vec,
            &self.trn_dir_vec,
            &self.classes_filepath,
            &self.label_filename,
        )
        .expect("init dataset failed")
    }
}

pub struct HWDBHdf5Dataset {
    hdf5_file: H5File,
    meta_list: Vec<(String, String, u16, u16)>,
    classes: Vec<String>,
    _dataset_type: HWDBHdf5DatasetType,
}

#[derive(Debug, Clone, Serialize, Deserialize, Copy)]
pub enum HWDBHdf5DatasetType {
    Train,
    Test,
    Valid,
    Unknown,
}

impl HWDBHdf5Dataset {
    pub fn from_hdf5(hdf5_path: &str, dataset_type: HWDBHdf5DatasetType) -> anyhow::Result<Self> {
        let file = H5File::open(hdf5_path)?;

        // 读取类别信息
        let classes_dataset: Dataset = file.dataset("classes")?;
        let classes: Vec<String> = classes_dataset
            .read_raw::<VarLenUnicode>()?
            .into_iter()
            .map(|v| v.to_string())
            .collect();

        // 读取图像路径和标签
        let tag_name = match dataset_type {
            HWDBHdf5DatasetType::Train => "trn_meta",
            HWDBHdf5DatasetType::Test => "tst_meta",
            HWDBHdf5DatasetType::Valid => "vld_meta",
            HWDBHdf5DatasetType::Unknown => "ukn_meta",
        };
        let dataset_meta: Dataset = file.dataset(tag_name)?;
        let meta_list = dataset_meta
            .read_raw::<ImageMeta>()?
            .into_iter()
            .map(|v| (v.name.to_string(), v.label.to_string(), v.width, v.height))
            .collect::<Vec<_>>();

        Ok(HWDBHdf5Dataset {
            hdf5_file: file,
            meta_list,
            classes,
            _dataset_type: dataset_type,
        })
    }
}

impl BurnDataset<ImageDatasetItem> for HWDBHdf5Dataset {
    fn get(&self, index: usize) -> Option<ImageDatasetItem> {
        if index >= self.meta_list.len() {
            return None;
        }

        let (name, label, width, height) = &self.meta_list[index];

        // 从HDF5文件中读取数据
        let image_dataset: Dataset = self.hdf5_file.dataset(name).ok()?;
        let image_data: Vec<u8> = image_dataset.read_raw().ok()?;

        // 将标签转换为索引
        let label_idx = self.classes.iter().position(|c| c == label)?;

        // 创建像素数据
        let pixel_data = image_data.into_iter().map(|b| PixelDepth::U8(b)).collect();

        Some(ImageDatasetItem {
            image: pixel_data,
            image_width: *width as usize,
            image_height: *height as usize,
            annotation: Annotation::Label(label_idx),
            image_path: name.clone(),
        })
    }

    fn len(&self) -> usize {
        self.meta_list.len()
    }
}

#[derive(Config, Debug)]
pub struct HWDBHdf5DatasetConfig {
    h5_filepath: String,
}

impl HWDBHdf5DatasetConfig {
    pub fn train(&self) -> HWDBHdf5Dataset {
        HWDBHdf5Dataset::from_hdf5(&self.h5_filepath, HWDBHdf5DatasetType::Train)
            .expect("init dataset failed")
    }

    pub fn test(&self) -> HWDBHdf5Dataset {
        HWDBHdf5Dataset::from_hdf5(&self.h5_filepath, HWDBHdf5DatasetType::Test)
            .expect("init dataset failed")
    }
}
